#!/usr/bin/env python3

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import cv2
import numpy as np
import serial
import time


class StereoBallDetector(Node):
    def __init__(self):
        super().__init__('stereo_ball_detector')
        
        # 协议常量
        self.FRAME_HEADER = 0x66
        self.FRAME_LENGTH = 0x07
        self.FRAME_TAIL = 0x5B
        self.SEND_INTERVAL = 0.05  # 发送间隔（20Hz）
        
        # 初始化相机参数
        self.cam_params = self.init_camera_params()
        self.screen_center = (self.cam_params["size"][0] // 2, self.cam_params["size"][1] // 2)
        self.get_logger().info(f"图像中心点: {self.screen_center}")
        
        # 初始化立体匹配器
        self.stereo, self.right_matcher, self.wls_filter = self.init_stereo_matcher(
            block_size=5, 
            num_disparities=96,
            speckle_size=100
        )
        
        # 初始化串口
        self.ser = self.init_serial()
        
        # 图像处理工具
        self.bridge = CvBridge()
        
        # 时间控制
        self.last_send_time = time.time()
        
        # 订阅图像话题
        self.image_sub = self.create_subscription(
            Image,
            '/image_raw',
            self.image_callback,
            10
        )
        
        # 创建显示窗口
        cv2.namedWindow("Rectified Color", cv2.WINDOW_NORMAL)
        cv2.namedWindow("Depth Map", cv2.WINDOW_NORMAL)
        
        self.get_logger().info("双目橙色球检测节点已启动（优化版）")

    def init_camera_params(self):
        """初始化双目相机参数（保持不变）"""
        left_camera_matrix = np.array([[715.088662856028, -1.27750597113329, 321.725929157012],
                                       [0, 717.938092436680, 251.352481762661],
                                       [0., 0., 1.]])
        right_camera_matrix = np.array([[712.537452801351, -1.49459241948369, 325.232106968860],
                                        [0, 715.552769066380, 250.629910089625],
                                        [0., 0., 1.]])

        left_distortion = np.array(
            [[-0.440986151974959, 0.327464329751247, -0.000547150225973162, -0.00126940700042557, -0.335074894313998]])
        right_distortion = np.array(
            [[-0.444122749137210, 0.348225581098782, 0.00176792559350367, 0.00123611244757822, -0.660993977721134]])

        R = np.array([[0.999866011965132, 0.000416637643297468, 0.0163641232584165],
                      [-0.000544744807817963, 0.999969236927380, 0.00782486114692518],
                      [-0.0163603597159969, -0.00783272698033651, 0.999835479975588]])
        T = np.array([-59.8387197927031, 0.0107452720401373, 0.330883446596301])
        size = (640, 480)

        R1, R2, P1, P2, Q, _, _ = cv2.stereoRectify(
            left_camera_matrix, left_distortion,
            right_camera_matrix, right_distortion,
            size, R, T
        )

        left_map1, left_map2 = cv2.initUndistortRectifyMap(
            left_camera_matrix, left_distortion, R1, P1, size, cv2.CV_16SC2)
        right_map1, right_map2 = cv2.initUndistortRectifyMap(
            right_camera_matrix, right_distortion, R2, P2, size, cv2.CV_16SC2)

        return {
            "left_camera_matrix": left_camera_matrix,
            "right_camera_matrix": right_camera_matrix,
            "left_distortion": left_distortion,
            "right_distortion": right_distortion,
            "R": R,
            "T": T,
            "size": size,
            "Q": Q,
            "left_map1": left_map1,
            "left_map2": left_map2,
            "right_map1": right_map1,
            "right_map2": right_map2,
            "baseline": abs(T[0])  # 基线（毫米）
        }

    def init_stereo_matcher(self, block_size=5, num_disparities=128, speckle_size=200):
        """初始化立体匹配器（保持不变）"""
        img_channels = 3
        stereo = cv2.StereoSGBM_create(
            minDisparity=0,
            numDisparities=num_disparities,
            blockSize=block_size,
            P1=8 * img_channels * block_size * block_size,
            P2=32 * img_channels * block_size * block_size,
            disp12MaxDiff=5,
            preFilterCap=63,
            uniquenessRatio=15,
            speckleWindowSize=speckle_size,
            speckleRange=32,
            mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY
        )

        right_matcher = cv2.ximgproc.createRightMatcher(stereo)
        wls_filter = cv2.ximgproc.createDisparityWLSFilter(matcher_left=stereo)
        wls_filter.setLambda(8000)
        wls_filter.setSigmaColor(1.5)

        return stereo, right_matcher, wls_filter

    def init_serial(self, port='/dev/ttyUSB0', baudrate=115200):
        """初始化串口通信（保持不变）"""
        try:
            ser = serial.Serial(
                port=port,
                baudrate=baudrate,
                bytesize=serial.EIGHTBITS,
                parity=serial.PARITY_NONE,
                stopbits=serial.STOPBITS_ONE,
                timeout=0.1
            )
            self.get_logger().info(f"串口已打开: {ser.port}")
            return ser
        except serial.SerialException as e:
            self.get_logger().error(f"串口打开失败: {e}")
            self.get_logger().warn("将以无串口模式运行")
            return None

    def rectify_image(self, img, map1, map2):
        """对图像进行畸变矫正（保持不变）"""
        return cv2.remap(img, map1, map2, cv2.INTER_LINEAR)

    def compute_disparity_and_depth(self, left_rect, right_rect):
        """计算视差图（移除全图三维计算）"""
        # 计算视差
        disparity_left = self.stereo.compute(left_rect, right_rect)
        disparity_right = self.right_matcher.compute(right_rect, left_rect)

        # 应用WLS滤波优化视差图
        filtered_disparity = self.wls_filter.filter(disparity_left, left_rect, None, disparity_right)
        filtered_disparity = np.clip(filtered_disparity, 0, None)

        # 视差图归一化（用于显示）
        disp_norm = cv2.normalize(filtered_disparity, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
        disp_color = cv2.applyColorMap(disp_norm, cv2.COLORMAP_JET)

        # 不再计算全图三维坐标，仅返回视差图
        return filtered_disparity, disp_norm, disp_color

    def detect_orange_ball(self, rect_color, filtered_disparity):
        """简化目标检测流程（优化颜色分割和轮廓处理）"""
        # 转换到HSV颜色空间
        hsv = cv2.cvtColor(rect_color, cv2.COLOR_BGR2HSV)
        lower_orange = np.array([0, 130, 130])
        upper_orange = np.array([179, 255, 255])
        
        # 颜色分割
        mask = cv2.inRange(hsv, lower_orange, upper_orange)

        # 优化形态学操作
        kernel = np.ones((3, 3), np.uint8)  # 原5x5 -> 3x3
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=1)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1)
        
        #cv2.imshow("mask",mask)

        # 查找轮廓
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        if not contours:
            return None
        
        # 提前过滤小面积轮廓（在找最大轮廓前过滤，减少计算）
        min_area = 50  # 可根据实际情况调整
        valid_contours = [cnt for cnt in contours if cv2.contourArea(cnt) > min_area]
        if not valid_contours:
            return None
            
        # 筛选面积最大的轮廓
        max_contour = max(valid_contours, key=cv2.contourArea)
        area = cv2.contourArea(max_contour)

        # 计算轮廓中心
        (cx_rect, cy_rect), radius = cv2.minEnclosingCircle(max_contour)
        cx_rect, cy_rect = int(cx_rect), int(cy_rect)

        # 仅在目标位置计算深度（替代全图三维计算）
        if 0 <= cx_rect < filtered_disparity.shape[1] and 0 <= cy_rect < filtered_disparity.shape[0]:
            # 获取目标位置的视差值（注意：视差图需转换为浮点数）
            disp_val = filtered_disparity[cy_rect, cx_rect] / 16.0  # SGBM输出视差需除以16
            
            if disp_val <= 0:  # 无效视差
                return None

            # 计算实际距离（米）：距离 = (基线 × 焦距) / 视差值
            # 基线单位转换为米（原单位毫米），焦距取左相机fx
            baseline_m = self.cam_params["baseline"] / 1000.0
            fx = self.cam_params["left_camera_matrix"][0, 0]  # 焦距（像素）
            distance = (baseline_m * fx) / disp_val  # 核心公式

            return {
                "rect_center": (cx_rect, cy_rect),
                "distance": distance,
                "contour": max_contour,
                "area": area
            }

        return None

    def send_coordinate(self, x, y):
        """按协议发送坐标（保持不变）"""
        if not self.ser or not self.ser.is_open:
            return

        x_clamped = np.clip(int(x), 0, 65535)
        y_clamped = np.clip(int(y), 0, 65535)

        x_low = x_clamped & 0xFF
        x_high = (x_clamped >> 8) & 0xFF
        y_low = y_clamped & 0xFF
        y_high = (y_clamped >> 8) & 0xFF

        frame = bytearray([
            self.FRAME_HEADER, self.FRAME_LENGTH,
            x_low, x_high, y_low, y_high,
            self.FRAME_TAIL
        ])

        try:
            self.ser.write(frame)
            self.get_logger().debug(f"发送坐标: ({x_clamped}, {y_clamped})")
        except serial.SerialException as e:
            self.get_logger().error(f"串口发送失败: {e}")

    def image_callback(self, msg):
        """处理订阅的图像消息（修改三维计算相关部分）"""
        try:
            cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
        except Exception as e:
            self.get_logger().error(f"图像转换失败: {e}")
            return

        # 分割左右相机图像
        frame_left = cv_image[:, :640]
        frame_right = cv_image[:, 640:]

        # 转换为灰度图（用于立体匹配）
        gray_left = cv2.cvtColor(frame_left, cv2.COLOR_BGR2GRAY)
        gray_right = cv2.cvtColor(frame_right, cv2.COLOR_BGR2GRAY)

        # 畸变矫正
        left_rect_gray = self.rectify_image(gray_left, self.cam_params["left_map1"], self.cam_params["left_map2"])
        left_rect_color = self.rectify_image(frame_left, self.cam_params["left_map1"], self.cam_params["left_map2"])
        right_rect_gray = self.rectify_image(gray_right, self.cam_params["right_map1"], self.cam_params["right_map2"])

        # 计算视差（不再返回全图三维坐标）
        filtered_disparity, _, disp_color = self.compute_disparity_and_depth(left_rect_gray, right_rect_gray)

        # 检测橙色球（传入视差图而非三维坐标）
        ball_info = self.detect_orange_ball(left_rect_color, filtered_disparity)

        # 绘制检测结果（保持不变）
        target_x, target_y = self.screen_center
        if ball_info:
            cx_rect, cy_rect = ball_info["rect_center"]
            distance = ball_info["distance"]
            contour = ball_info["contour"]
            target_x, target_y = cx_rect, cy_rect

            cv2.drawContours(left_rect_color, [contour], -1, (0, 255, 0), 2)
            cv2.circle(left_rect_color, (cx_rect, cy_rect), 5, (0, 0, 255), -1)
            cv2.line(left_rect_color, self.screen_center, (cx_rect, cy_rect), (0, 255, 255), 2)
            cv2.putText(left_rect_color, f"Center: ({cx_rect}, {cy_rect})", (10, 30),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
            cv2.putText(left_rect_color, f"Distance: {distance:.3f}m", (10, 60),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
        else:
            cv2.putText(left_rect_color, "No Target", (10, 30),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)

        # 绘制图像中心点（保持不变）
        cv2.circle(left_rect_color, self.screen_center, 5, (255, 0, 0), -1)
        cv2.line(left_rect_color, (self.screen_center[0]-10, self.screen_center[1]),
                (self.screen_center[0]+10, self.screen_center[1]), (255, 0, 0), 2)
        cv2.line(left_rect_color, (self.screen_center[0], self.screen_center[1]-10),
                (self.screen_center[0], self.screen_center[1]+10), (255, 0, 0), 2)

        # 发送坐标（保持不变）
        current_time = time.time()
        if current_time - self.last_send_time > self.SEND_INTERVAL:
            self.send_coordinate(target_x, target_y)
            self.last_send_time = current_time

        # 显示图像（保持不变）
        cv2.imshow("Rectified Color", left_rect_color)
        cv2.imshow("Depth Map", disp_color)
        cv2.waitKey(1)


def main(args=None):
    rclpy.init(args=args)
    detector = StereoBallDetector()
    try:
        rclpy.spin(detector)
    except KeyboardInterrupt:
        pass
    finally:
        cv2.destroyAllWindows()
        if detector.ser and detector.ser.is_open:
            detector.ser.close()
            detector.get_logger().info("串口已关闭")
        detector.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
